from collections import defaultdict
from datetime import datetime
from pymongo import UpdateOne
from conn import ac_task_col, clip_annotated_record, check_frame_label_col
from tqdm import tqdm
from loguru import logger


def remove_repeat_clip(anno_type: int):
    """
    删除 ac task 中重复的 clip, 保留最近加入的 clip
    """
    task_list = (
        ac_task_col.find(
            {"annotate_type": anno_type},
            {"clip_ids": 1, "clip_annotated_records": 1, "created_at": 1},
        )
        .sort([("_id", -1)])
        .to_list()
    )
    clip_id_set = set()
    for task in tqdm(task_list):
        repeat_clip_ids = []
        for clip_id in task["clip_ids"]:
            if clip_id in clip_id_set:
                repeat_clip_ids.append(clip_id)
            clip_id_set.add(clip_id)
        if len(repeat_clip_ids) == 0:
            continue
        # 当前任务中应该删除的重复 clip
        logger.warning(
            f"{task['_id']} repeat_clip_ids: {len(repeat_clip_ids)}, {task['created_at']}"
        )
        logger.warning(f"repeat_clip_ids: {repeat_clip_ids}")
        clip_info_map = {
            clip["clip_id"]: clip for clip in task["clip_annotated_records"]
        }
        new_clip_ids, new_clip_infos = [], []
        for clip_id in task["clip_ids"]:
            if clip_id in repeat_clip_ids:
                continue
            new_clip_ids.append(clip_id)
            new_clip_infos.append(clip_info_map[clip_id])
        if len(new_clip_ids) == 0:
            logger.info(f"task {task['_id']} is empty, should remove")
            ac_task_col.delete_one({"_id": task["_id"]})
        else:
            logger.info(
                f"task {task['_id']} has {len(new_clip_ids)}/{len(task['clip_ids'])} clips"
            )
            ac_task_col.update_one(
                {"_id": task["_id"]},
                {
                    "$set": {
                        "clip_ids": new_clip_ids,
                        "clip_ids_num": len(new_clip_ids),
                        "clip_annotated_records": new_clip_infos,
                    }
                },
            )


def remove_not_confirmed_clip(anno_type: int):
    """
    移除 ac task 中未确认的 clip, 一般是被重新驳回导致的
    """
    task_list = (
        ac_task_col.find(
            {"annotate_type": anno_type},
            {"clip_ids": 1, "clip_annotated_records": 1, "created_at": 1},
        )
        .sort([("_id", -1)])
        .to_list()
    )
    for task in tqdm(task_list):
        clip_list = clip_annotated_record.find(
            {"clip_id": {"$in": task["clip_ids"]}, "annotate_type": anno_type},
            {"state": 1, "clip_id": 1},
        ).to_list()

        not_confirmed_clip_ids = []
        for clip in clip_list:
            if clip["state"] != 2:
                not_confirmed_clip_ids.append(clip["clip_id"])
        if len(not_confirmed_clip_ids) == 0:
            continue

        logger.warning(
            f"{task['_id']} not_confirmed_clip_ids: {len(not_confirmed_clip_ids)}, {task['created_at']}"
        )

        clip_info_map = {
            clip["clip_id"]: clip for clip in task["clip_annotated_records"]
        }
        new_clip_ids, new_clip_infos = [], []
        for clip_id in task["clip_ids"]:
            if clip_id in not_confirmed_clip_ids:
                continue
            new_clip_ids.append(clip_id)
            new_clip_infos.append(clip_info_map[clip_id])
        if len(new_clip_ids) == 0:
            logger.info(f"task {task['_id']} is empty, should remove")
            ac_task_col.delete_one({"_id": task["_id"]})
        else:
            logger.info(
                f"task {task['_id']} has {len(new_clip_ids)}/{len(task['clip_ids'])} clips"
            )
            ac_task_col.update_one(
                {"_id": task["_id"]},
                {
                    "$set": {
                        "clip_ids": new_clip_ids,
                        "clip_ids_num": len(new_clip_ids),
                        "clip_annotated_records": new_clip_infos,
                    }
                },
            )


def repair_clip_ac_status(anno_type: int):
    """
    修复 ac task 中 clip 的验收状态，主要是修复 ac 状态为 10 或者 20 的 clip，其验收状态为 20 的情况
    修复逻辑：
    1. 找到所有 ac 状态为 10 或者 20 且确认的 clip
    2. 找到这些 clip 对应的 frame_check 记录，如果 frame_check 记录数量大于等于 19 或者 190，则更新 ac 状态为 20，否则为 10
    3. 更新 ac task 中 clip 的验收状态
    """
    task_list = (
        ac_task_col.find(
            {"annotate_type": anno_type},
            {"clip_ids": 1, "status": 1, "frequency": 1},
        )
        .sort([("_id", -1)])
        .to_list()
    )
    for task in tqdm(task_list):
        clip_ids = task["clip_ids"]
        clip_records = clip_annotated_record.find(
            {
                "clip_id": {"$in": clip_ids},
                "annotate_type": anno_type,
                "state": 2,  # 确认状态，才尝试修复其验收状态，非确认状态，暂时不关注
            },
            {"accept_status": 1, "clip_id": 1, "state": 1, "accept_finished_at": 1},
        ).to_list()
        # 记录每个 clip 的更新操作，应该更新为 10 还是 20，需要根据 frame_check 表中的数量进行判断
        update_ops = []
        for clip in clip_records:
            clip_id = clip["clip_id"]
            cur_ac_status = clip.get("accept_status", None)
            cur_state = clip.get("state", None)
            cur_ac_finished_at = clip.get("accept_finished_at", None)
            # clip 状态确认，且 ac 状态为 10 或 (20 + 存在 accept_finished_at)，则跳过
            if cur_ac_status == 10:
                continue
            if cur_ac_status == 20 and cur_ac_finished_at is not None:
                continue
            logger.warning(
                f"task {task['_id']} clip {clip_id} status error: state {cur_state} accept_status {cur_ac_status} ac_time: {cur_ac_finished_at}"
            )
            ac_status, ac_time = None, None  # 非确认状态，都为 null 值
            if cur_state == 2:
                # 确认状态，则需要更新成 10 或者 20, frame_check 表中有指定数量的记录，则更新成 20 和具体的时间
                # TODO 这里其实有点问题，可能是二次确认，需要二次加入任务中进行验收，但这里没有考虑这种情况，直接标记完成了
                ac_status, ac_time = 10, None
                latest_check_frames = (
                    check_frame_label_col.find(
                        dict(clip_id=clip_id, annotate_type=anno_type),
                        {"created_at": 1},
                    )
                    .sort([("_id", -1)])
                    .to_list()
                )
                if (task["frequency"] == "1Hz" and len(latest_check_frames) >= 19) or (
                    task["frequency"] == "10Hz" and len(latest_check_frames) >= 190
                ):
                    ac_status = 20
                    ac_time = latest_check_frames[0]["created_at"]
            # 都相等的话，无需更新
            if ac_status == cur_ac_status and ac_time == cur_ac_finished_at:
                continue
            logger.info(
                f"update task {task['_id']} clip {clip_id} status: accept_status {ac_status}, ac_time {ac_time}"
            )
            update_ops.append(
                UpdateOne(
                    {"clip_id": clip_id, "annotate_type": anno_type},
                    {
                        "$set": {
                            "accept_status": ac_status,
                            "accept_finished_at": ac_time,
                        }
                    },
                )
            )
        if update_ops:
            clip_annotated_record.bulk_write(update_ops, ordered=False)


# 时间相差 10s 内，认为是同一个时间
def same_time(t1: datetime | None, t2: datetime | None, threshold: int = 10):
    if t1 is None and t2 is None:
        return True
    if t1 is None or t2 is None:
        return False
    return abs((t1 - t2).total_seconds()) < threshold


def repair_ac_task_status(anno_type: int):
    """
    修复 ac task 的状态，通过对不同 ac 状态的 clip 进行计数判断，来决定其最终状态
    """
    task_list = (
        ac_task_col.find(
            {"annotate_type": anno_type},
            {"clip_ids": 1, "status": 1, "frequency": 1, "finished_at": 1},
        )
        .sort([("_id", -1)])
        .to_list()
    )
    update_ops = []
    for task in tqdm(task_list):
        # 检查 task 中每个 clip 的 accept_status
        clip_ids = task["clip_ids"]
        clip_records = (
            clip_annotated_record.find(
                {
                    "clip_id": {"$in": clip_ids},
                    "annotate_type": anno_type,
                    "state": 2,  # 确认状态，才尝试修复其验收状态，非确认状态，暂时不关注
                },
                {"accept_status": 1, "clip_id": 1, "state": 1, "accept_finished_at": 1},
            )
            .sort([("accept_finished_at", -1)])
            .to_list()
        )

        # clip 验收状态计数
        clip_status_cnt = defaultdict(int)
        for clip in clip_records:
            clip_status_cnt[clip.get("accept_status")] += 1

        new_status, new_finished_at = None, None
        if clip_status_cnt[20] == len(clip_ids):
            new_status = 3  # 完成
            new_finished_at = (
                task.get("finished_at") or clip_records[0]["accept_finished_at"]
            )
        elif clip_status_cnt[20] > 0:
            new_status = 2  # 进行中
            new_finished_at = None
        else:
            new_status = 1  # 待验收
            new_finished_at = None
        if new_status == task["status"] and same_time(
            task.get("finished_at"), new_finished_at, 60
        ):
            # 状态相同，不再更新
            continue
        logger.info(
            f"task {task['_id']} status error: {task['status']} -> {new_status}, finished_at: {task.get('finished_at')} -> {new_finished_at}, clip_cnt: {len(clip_ids)}, {clip_status_cnt}"
        )
        update_ops.append(
            UpdateOne(
                {"_id": task["_id"]},
                {
                    "$set": {
                        "status": new_status,
                        "finished_at": new_finished_at,
                    }
                },
            )
        )
        if len(update_ops) > 100:
            ac_task_col.bulk_write(update_ops, ordered=False)
            update_ops = []

    if update_ops:
        ac_task_col.bulk_write(update_ops, ordered=False)


annotate_types = [2, 19, 20, 21, 29, 31]


def main():
    # for anno_type in annotate_types:
    #     logger.info(f"handle repeat clips for anno_type: {anno_type}")
    #     remove_repeat_clip(anno_type)

    # for anno_type in annotate_types:
    #     logger.info(f"handle not comfirmed clips for anno_type: {anno_type}")
    #     remove_not_confirmed_clip(anno_type)

    # for anno_type in annotate_types:
    #     logger.info(f"repair clip ac status for anno_type: {anno_type}")
    #     repair_clip_ac_status(anno_type)

    # for anno_type in annotate_types:
    #     logger.info(f"repair ac task status for anno_type: {anno_type}")
    #     repair_ac_task_status(anno_type)
    pass


if __name__ == "__main__":
    main()
